.globl relu

.text
# ==============================================================================
# FUNCTION: Performs an inplace element-wise ReLU on an array of ints
# Arguments:
# 	a0 (int*) is the pointer to the array
#	a1 (int)  is the # of elements in the array
# Returns:
#	None
#
# If the length of the vector is less than 1, 
# this function exits with error code 8.
# ==============================================================================
relu:
    # Prologue
    addi sp, sp, -24
    sw ra, 0(sp)
    sw s0, 4(sp)
    sw s1, 8(sp)
    sw s2, 12(sp)
    sw s3, 16(sp)
    sw s4, 20(sp)

    # save arguments
    add s0, a0, x0
    add s1, a1, x0

    # initial
    add s2, x0, x0 # s2 = 0
    add s3, x0, x0 # s3 = 0
    # if a1 >= 1, goto loop_start
    addi s4, x0, 1 # s4 = 1
    bge a1, s4, loop_start
    # call exit, status code = 8
    addi a1, x0, 8
    jal ra, exit2


loop_start:
    # if s2 == a1, goto loop_end
    beq s2, s1, loop_end
    # if a0[s2] < 0, goto loop_continue
    slli s3, s2, 2 # s3 = s2 * 4
    add s3, s0, s3
    lw s4, 0(s3) # s4 = a0[s2]
    blt s4, x0, loop_continue
    # t0++
    addi s2, s2, 1
    # goto loop_start
    j loop_start


loop_continue:
    # a0[s2] = 0
    sw x0, 0(s3)
    # s2++
    addi s2, s2, 1
    # goto loop_start
    j loop_start


loop_end:
    # Epilogue
    lw ra, 0(sp)
    lw s0, 4(sp)
    lw s1, 8(sp)
    lw s2, 12(sp)
    lw s3, 16(sp)
    lw s4, 20(sp)
    addi sp, sp, 24

    ret
